# Load required packages
librarian::shelf(ggthemes, tibble, quiet = TRUE)

# Set PRL theme
theme_prl <- function(base_size = 14, base_family = "sans") {
  colors <- deframe(ggthemes_data[["fivethirtyeight"]])
  (theme_foundation(base_size = base_size, base_family = base_family)
  + theme(
      line = element_line(colour = "black"),
      rect = element_rect(
        linetype = 0, colour = NA
      ),
      text = element_text(colour = colors["Dark Gray"]),
      axis.title = element_text(size = rel(.8)),
      axis.text = element_text(color = colors["Dark Gray"], face = "bold"),
      # axis.ticks = element_blank(),
      axis.line = element_blank(),
      legend.background = element_rect(),
      legend.title = element_text(hjust = 0, size = rel(.7), face = "bold"),
      legend.text = element_text(hjust = 0, size = rel(.7)),
      legend.position = "bottom",
      legend.direction = "horizontal",
      legend.box = "vertical",
      panel.grid = element_line(colour = NULL),
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      plot.title = element_text(hjust = 0, size = rel(1.5), face = "bold"),
      plot.margin = unit(c(.5, .5, .5, .5), "lines"),
      strip.text = element_text(face = "bold"),
      plot.caption = element_text(),
      strip.background = element_rect()
    )
  )
}

# Set FiveThirtyEight functions (needed for PRL theme)
fivethirtyeight_pal <- function() {
  colors <- deframe(ggthemes_data[["fivethirtyeight"]])
  values <- unname(colors[c("Blue", "Red", "Green")])
  max_n <- length(values)
  f <- manual_pal(values)
  attr(f, "max_n") <- max_n
  f
}

scale_colour_fivethirtyeight <- function(...) {
  discrete_scale("colour", "economist", fivethirtyeight_pal(), ...)
}

scale_color_fivethirtyeight <- scale_colour_fivethirtyeight

scale_fill_fivethirtyeight <- function(...) {
  discrete_scale("fill", "economist", fivethirtyeight_pal(), ...)
}

# Set up section functions
GeomSection <- ggproto("GeomSection", GeomPolygon,
  default_aes = list(
    fill = "gray", size = 0, alpha = 0.2, colour = NA, linetype = "dashed"
  ),
  required_aes = c("slope", "intercept", "above"),
  draw_panel = function(data, panel_params, coord) {
    ranges <- coord$backtransform_range(panel_params)
    data$group <- seq_len(nrow(data))
    data <- data %>%
      group_by_all() %>%
      do(buildPoly(.$slope, .$intercept, .$above, ranges$x, ranges$y)) %>%
      unnest()
    GeomPolygon$draw_panel(data, panel_params, coord)
  }
)

geom_section <- function(mapping = NULL, data = NULL, ..., slope, intercept, above,
                         na.rm = FALSE, show.legend = NA) {
  if (missing(mapping) && missing(slope) && missing(intercept) && missing(above)) {
    slope <- 1
    intercept <- 0
    above <- TRUE
  }
  if (!missing(slope) || !missing(intercept) || !missing(above)) {
    if (missing(slope)) {
      slope <- 1
    }
    if (missing(intercept)) {
      intercept <- 0
    }
    if (missing(above)) {
      above <- TRUE
    }
    data <- data.frame(intercept = intercept, slope = slope, above = above)
    mapping <- aes(intercept = intercept, slope = slope, above = above)
    show.legend <- FALSE
  }
  layer(
    data = data, mapping = mapping, stat = StatIdentity,
    geom = GeomSection, position = PositionIdentity, show.legend = show.legend,
    inherit.aes = FALSE, params = list(na.rm = na.rm, ...)
  )
}

buildPoly <- function(slope, intercept, above, xr, yr) {
  yCross <- (yr - intercept) / slope
  xCross <- (slope * xr) + intercept

  # Build polygon by cases
  if (above & (slope >= 0)) {
    rs <- data.frame(x = -Inf, y = Inf)
    if (xCross[1] < yr[1]) {
      rs <- rbind(rs, c(-Inf, -Inf), c(yCross[1], -Inf))
    } else {
      rs <- rbind(rs, c(-Inf, xCross[1]))
    }
    if (xCross[2] < yr[2]) {
      rs <- rbind(rs, c(Inf, xCross[2]), c(Inf, Inf))
    } else {
      rs <- rbind(rs, c(yCross[2], Inf))
    }
  }
  if (!above & (slope >= 0)) {
    rs <- data.frame(x = Inf, y = -Inf)
    if (xCross[1] > yr[1]) {
      rs <- rbind(rs, c(-Inf, -Inf), c(-Inf, xCross[1]))
    } else {
      rs <- rbind(rs, c(yCross[1], -Inf))
    }
    if (xCross[2] > yr[2]) {
      rs <- rbind(rs, c(yCross[2], Inf), c(Inf, Inf))
    } else {
      rs <- rbind(rs, c(Inf, xCross[2]))
    }
  }
  if (above & (slope < 0)) {
    rs <- data.frame(x = Inf, y = Inf)
    if (xCross[1] < yr[2]) {
      rs <- rbind(rs, c(-Inf, Inf), c(-Inf, xCross[1]))
    } else {
      rs <- rbind(rs, c(yCross[2], Inf))
    }
    if (xCross[2] < yr[1]) {
      rs <- rbind(rs, c(yCross[1], -Inf), c(Inf, -Inf))
    } else {
      rs <- rbind(rs, c(Inf, xCross[2]))
    }
  }
  if (!above & (slope < 0)) {
    rs <- data.frame(x = -Inf, y = -Inf)
    if (xCross[1] > yr[2]) {
      rs <- rbind(rs, c(-Inf, Inf), c(yCross[2], Inf))
    } else {
      rs <- rbind(rs, c(-Inf, xCross[1]))
    }
    if (xCross[2] > yr[1]) {
      rs <- rbind(rs, c(Inf, xCross[2]), c(Inf, -Inf))
    } else {
      rs <- rbind(rs, c(yCross[1], -Inf))
    }
  }
  return(rs)
}
